import os, time, datetime
import gym
from gym import spaces
from gym.utils import seeding
import numpy as np
import pybullet as p
import cv2
# from keras.models import load_model

from .world_creation import WorldCreation


class RobotEnv(gym.Env):
    def __init__(self, robot_type='jaco', frame_skip=5, time_step=0.02, action_robot_len=12, obs_robot_len=30):
        # Start the bullet physics server
        self.id = p.connect(p.DIRECT)
        self.id = 0
        # print('Physics server ID:', self.id)
        self.gui = False
        self.robot_type = robot_type
        ## gym

        self.action_robot_len = action_robot_len
        self.obs_robot_len = obs_robot_len
        self.action_space = spaces.Box(low=np.array([-1.0]*(self.action_robot_len)), high=np.array([1.0]*(self.action_robot_len)), dtype=np.float32)
        self.observation_space = spaces.Box(low=np.array([-1.0]*(self.obs_robot_len)), high=np.array([1.0]*(self.obs_robot_len)), dtype=np.float32)

        # Execute actions at 10 Hz by default. A new action every 0.1 seconds
        self.frame_skip = frame_skip
        self.time_step = time_step
        self.seed(5)

        self.world_creation = WorldCreation(self.id, robot_type=robot_type, time_step=self.time_step, np_random=self.np_random)

        self._render_width = 50
        self._render_height = 50

        self.width = 1920
        self.height = 1080

    def seed(self, seed=None):
        self.np_random, seed = seeding.np_random(seed)
        return [seed]
    ## for reaching
    def step(self, action):
        raise NotImplementedError('Implement observations')

    def _get_obs(self, forces):
        raise NotImplementedError('Implement observations')

    def reset(self):
        raise NotImplementedError('Implement reset')

    def take_step(self, action,  gains=0.1,  step_sim=True):

        action = np.clip(action, a_min=self.action_space.low, a_max=self.action_space.high)


        action *= (1/self.frame_skip)

        action_robot = action
        indices = [2,3,4,5,6,7,9,10,11,12,13,14]
        robot_joint_states = p.getJointStates(self.robot, jointIndices=indices, physicsClientId=self.id)
        robot_joint_positions = np.array([x[0] for x in robot_joint_states])

        for _ in range(self.frame_skip):
            action_robot[robot_joint_positions + action_robot < self.robot_lower_limits] = 0
            action_robot[robot_joint_positions + action_robot > self.robot_upper_limits] = 0
            robot_joint_positions += action_robot

        p.setJointMotorControlArray(self.robot, jointIndices=indices, controlMode=p.POSITION_CONTROL, targetPositions=robot_joint_positions, positionGains=np.array([gains]*self.action_robot_len), forces=[40,80,40,20,20,20,2,2,2,2,2,2], physicsClientId=self.id)


        if step_sim:
            for _ in range(self.frame_skip):
                p.stepSimulation(physicsClientId=self.id)

                
    ## 重置关节 应用于reset()
    def position_robot_toc(self, robot, joints, joint_indices, lower_limits, upper_limits, pos_offset=np.zeros(3)):

        if type(joints) == int:
            joints = [12]
            joint_indices = [[2,3,4,5,6,7,9,10,11,12,13,14]]
            lower_limits = np.append(lower_limits[2:8],lower_limits[9:])
            upper_limits = np.append(upper_limits[2:8],upper_limits[9:])

        best_position = np.array([0, 0, 0])
        best_orientation = np.array([0.0, 0.0, 0,1])
        best_start_joint_poses = [[0,2.9,1.3,-2.07,1.4,0,0,0,0,0,0,0]]

        # spawn robot
        p.resetBasePositionAndOrientation(robot, pos_offset + best_position, best_orientation, physicsClientId=self.id)
        ## 设置 机器人的关节
        for i in range(12):
            p.resetJointState(robot, jointIndex=joint_indices[0][i], targetValue=best_start_joint_poses[0][i], targetVelocity=0, physicsClientId=self.id)
            
        
        return best_position, best_orientation, best_start_joint_poses

    def render(self, mode='human'):
        if mode =="rgb_array":
            novice_image = self.different_viewpoints(base_pos=[0,0,2],cam_dist=2,cam_yaw=0,cam_pitch=180,cam_roll=0)
            expert_image = self.different_viewpoints(base_pos=[0,0,1],cam_dist=2,cam_yaw=90,cam_pitch=-30,cam_roll=180)
            noviceImage = np.array(novice_image)[:, :, :3]
            expertImage = np.array(expert_image)[:, :, :3]

            return expertImage,noviceImage

        if not self.gui:
            self.gui = True
            p.disconnect(self.id)
            self.id = p.connect(p.GUI, options='--background_color_red=0.8 --background_color_green=0.9 --background_color_blue=1.0 --width=%d --height=%d' % (self.width, self.height))

            self.world_creation = WorldCreation(self.id, robot_type=self.robot_type, time_step=self.time_step, np_random=self.np_random)


